import cv2
import numpy as np
import math
import os

def make_chessboard(sqr_px=50, int_corners=(8,5), border_squares=2):
    """
    生成一张棋盘格图：
    - internal corners = int_corners (8×5)
    - 方格数 = int_corners + 1 = (9×6)
    - 每格边长 sqr_px
    - 四周留 border_squares 格的白边
    """
    ic, ir = int_corners      # 内部交点数：列, 行
    cols = ic + 1             # 方格列数
    rows = ir + 1             # 方格行数
    w0, h0 = cols * sqr_px, rows * sqr_px

    # 1) 生成棋盘本体
    board = np.zeros((h0, w0), np.uint8)
    for y in range(rows):
        for x in range(cols):
            if (x + y) % 2 == 0:
                board[y*sqr_px:(y+1)*sqr_px,
                      x*sqr_px:(x+1)*sqr_px] = 255
    board = cv2.cvtColor(board, cv2.COLOR_GRAY2BGR)

    # 2) 在四周加白边
    pad = border_squares * sqr_px
    board = cv2.copyMakeBorder(
        board, pad, pad, pad, pad,
        borderType=cv2.BORDER_CONSTANT,
        value=(255,255,255)
    )
    return board

def warp_chessboard_full(img, int_corners, pitch, yaw,
                         f=800, distance=1000, pad_sqr=2):
    """
    将 img 做透视变换，并放到一个动态大小的白底画布上：
    - int_corners: (cols,rows) 内部交点数量
    - pitch, yaw: 俯仰/偏航 角度 (°)
    - f: 焦距 (px)
    - distance: 棋盘平面到相机的距离
    - pad_sqr: 额外留几格白边防止贴边
    """
    h, w = img.shape[:2]
    # 原图四角 3D 坐标
    pts3d = np.array([[0,0,0], [w,0,0], [w,h,0], [0,h,0]], np.float32)

    # 相机内参
    cx, cy = w/2, h/2

    # 构造旋转矩阵 Rx, Ry
    ax = math.radians(pitch)
    ay = math.radians(yaw)
    Rx = np.array([[1,0,0],
                   [0,math.cos(ax),-math.sin(ax)],
                   [0,math.sin(ax), math.cos(ax)]], np.float32)
    Ry = np.array([[ math.cos(ay),0,math.sin(ay)],
                   [0,            1,           0],
                   [-math.sin(ay),0,math.cos(ay)]], np.float32)
    R = Ry.dot(Rx)
    t = np.array([[0],[0],[distance]], np.float32)

    # 投影四角到像素平面
    dst = []
    for X, Y, Z in pts3d:
        P = R.dot(np.array([[X],[Y],[Z]],np.float32)) + t
        x = (P[0]/P[2])*f + cx
        y = (P[1]/P[2])*f + cy
        dst.append([x, y])
    dst = np.array(dst, np.float32)

    # 单应 H: src->dst
    src = np.array([[0,0],[w,0],[w,h],[0,h]], np.float32)
    H, _ = cv2.findHomography(src, dst)

    # 计算输出画布范围，并加 pad_sqr 个格子的白边
    ic, ir = int_corners
    square_px = w / (ic + 1)
    pad = pad_sqr * square_px

    mn = dst.min(axis=0) - pad
    mx = dst.max(axis=0) + pad
    min_x, min_y = float(mn[0]), float(mn[1])
    max_x, max_y = float(mx[0]), float(mx[1])

    out_w = int(math.ceil(max_x - min_x))
    out_h = int(math.ceil(max_y - min_y))

    # 平移矩阵 T，将 (min_x,min_y) 映到 (0,0)
    T = np.array([[1, 0, -min_x],
                  [0, 1, -min_y],
                  [0, 0,      1]], np.float32)
    H2 = T.dot(H)

    # Warp 到新画布
    warped = cv2.warpPerspective(
        img, H2, (out_w, out_h),
        flags=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=(255,255,255)
    )
    return warped

def main():
    out_dir = "calib_images"
    os.makedirs(out_dir, exist_ok=True)

    # 1) 生成棋盘格
    square_px    = 50
    int_corners  = (8, 5)    # 内部交点数：8 列 × 5 行
    border_squares = 2       # 边上多留 2 格白边
    board = make_chessboard(square_px, int_corners, border_squares)

    # 2) 随机生成 15 组角度：±15°
    np.random.seed(42)
    pitches = np.random.uniform(-15, 15, 15)
    yaws    = np.random.uniform(-15, 15, 15)

    # 3) 逐张透视并保存
    for idx, (p, y) in enumerate(zip(pitches, yaws), start=1):
        warped = warp_chessboard_full(
            board, int_corners,
            pitch=p, yaw=y,
            f=800, distance=1000, pad_sqr=2
        )
        fname = os.path.join(out_dir, f"img{idx:02d}.jpg")
        cv2.imwrite(fname, warped)
        print(f"Saved {fname}: pitch={p:.1f}°, yaw={y:.1f}°")

    print("→ 15 张棋盘格标定图已生成到", out_dir)

if __name__ == "__main__":
    main()
